import os
import shutil
import uuid
from typing import Optional
from fastapi import FastAPI
from pydantic import BaseModel
from tqdm import tqdm
from preprocess.file_converter import convert_file
from preprocess.file_paser import parse_file
from preprocess.file_spliter import split_file, traverse_dir
from pasers.img_parser.img_ocr import image_ocr
from rapidocr import RapidOCR


# 定义请求体
class ProcessRequest(BaseModel):
    input: Optional[str] = "/preprocess_project/doc_paser-master/files"
    output: Optional[str] = ""
    chunk_size: int = 800
    chunk_overlap: int = 150
    use_cpu_only: bool = False
    keep_intermediate: bool = False


app = FastAPI(title="文档处理服务")


@app.post("/process")
def process_files(req: ProcessRequest):
    # 初始化 OCR
    use_gpu = not req.use_cpu_only
    ocr = RapidOCR(params={'EngineConfig.use_cuda': use_gpu, 'Global.with_onnx': True})

    # 默认输出目录
    if not req.output:
        req.output = os.path.dirname(req.input.rstrip('/'))

    SUPPORTED_FORMATS = {
        'doc', 'docx', 'xls', 'xlsx', 'pdf', 'ofd', 'xml', 'txt',
        'jpg', 'jpeg', 'png', 'jfif', 'exif', 'gif', 'tiff', 'bmp', 'wmf'
    }

    split_size = req.chunk_size
    register_dict = {}
    failed_ls = []

    base_name = os.path.basename(req.input.rstrip('/'))
    uuid_dir = os.path.join(req.output, f'{base_name}_uuid_files')
    os.makedirs(uuid_dir, exist_ok=True)

    namespace = uuid.NAMESPACE_DNS
    bar = tqdm(traverse_dir(req.input, None), desc='为文件准备uuid')
    for fp, fn in bar:
        file_suffix = fp.split('.')[-1]
        if file_suffix in SUPPORTED_FORMATS:
            uuid_name = str(uuid.uuid5(namespace, fp))
            register_dict[uuid_name] = {'path': fp, 'name': fn + '.' + file_suffix}
            shutil.copy(fp, os.path.join(uuid_dir, uuid_name + '.' + file_suffix))
        else:
            failed_ls.append({'filePath': fp, 'reason': '文件格式不支持'})

    converted_dir = os.path.join(req.output, f'{base_name}_converted')
    ofd_unzip_dir = os.path.join(req.output, f'{base_name}_ofd_unzipped')
    image_saved_dir = os.path.join(req.output, f'{base_name}_images')
    parse_out_dir = os.path.join(req.output, f'{base_name}_out')
    seg_out_dir = os.path.join(req.output, f'{base_name}_seg')

    # 执行原逻辑
    convert_file(failed_ls, uuid_dir, converted_dir, converted_to_pdf=False)
    parse_file(failed_ls, converted_dir, parse_out_dir, ofd_unzip_dir, image_saved_dir)
    image_ocr(parse_out_dir, image_saved_dir, ocr)
    split_file(
        failed_ls,
        parse_out_dir,
        converted_dir,
        seg_out_dir,
        split_size,
        req.chunk_overlap,
        register_dict,
        ocr,
        with_position=False
    )

    # 清理临时文件
    if not req.keep_intermediate:
        shutil.rmtree(uuid_dir, ignore_errors=True)
        shutil.rmtree(converted_dir, ignore_errors=True)
        shutil.rmtree(ofd_unzip_dir, ignore_errors=True)
        shutil.rmtree(parse_out_dir, ignore_errors=True)

    return {
        "register_dict": register_dict,
        "failed_files": failed_ls,
        "output_dir": req.output
    }
